"""
自定义Dataset类实现内存映射加载
多线程数据加载（num_workers=4）
支持GPU加速的张量运算

关键PyTorch特性应用：
张量核心优化：通过torch.cuda.amp实现混合精度训练
内存优化：使用pin_memory=True加速GPU数据传输
概率预测：通过MC Dropout实现不确定性估计
"""
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, TensorDataset
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt


# 生成模拟数据（保持与之前相同）
def generate_15min_data(days=30):
    timestamps = pd.date_range(start="2023-01-01", periods=days * 96, freq="15T")
    # 基础价格模式
    base_price = 50 + 20 * np.sin(2 * np.pi * np.arange(len(timestamps)) / 96)  # 日周期
    # 添加周周期
    week_cycle = 10 * np.sin(2 * np.pi * np.arange(len(timestamps)) / (96 * 7))
    # 实时特征
    demand = np.random.normal(5000, 500, len(timestamps))
    temperature = 15 + 10 * np.sin(2 * np.pi * np.arange(len(timestamps)) / 96)  # 与电价正相关
    # 构建DataFrame
    data = pd.DataFrame({
        'datetime': timestamps,
        'price': base_price + week_cycle + 0.1 * (demand - 5000) + 0.5 * temperature + np.random.normal(0, 3,
                                                                                                        len(timestamps)),
        'demand': demand,
        'temperature': temperature,
        'renewable_ratio': np.clip(np.random.beta(2, 5, len(timestamps)), 0, 0.6)  # 可再生能源占比
    })
    # 添加周末特征
    data['is_weekend'] = data['datetime'].dt.dayofweek.isin([5, 6]).astype(int)

    return data


# 数据预处理
def preprocess_data(data):
    # 时间特征工程
    data['15min_block'] = (data['datetime'].dt.hour * 4) + (data['datetime'].dt.minute // 15)
    data['day_of_week'] = data['datetime'].dt.dayofweek
    data['month'] = data['datetime'].dt.month
    # 滞后特征（使用前三天同一时段价格）
    for lag in [1 * 96, 2 * 96, 3 * 96]:  # 1天、2天、3天前
        data[f'price_lag_{lag}'] = data['price'].shift(lag)
    # 移动平均特征
    data['price_7day_ma'] = data['price'].rolling(window=7 * 96).mean()
    data = data.dropna()
    return data

# 数据预处理（改进为PyTorch兼容格式）
class EnergyDataset(Dataset):
    def __init__(self, data, look_back=672, pred_steps=96):
        self.scaler = MinMaxScaler(feature_range=(-1, 1))
        scaled_data = self.scaler.fit_transform(data)

        self.X = []
        self.y = []
        for i in range(len(scaled_data) - look_back - pred_steps):
            self.X.append(scaled_data[i:i + look_back])
            self.y.append(scaled_data[i + look_back:i + look_back + pred_steps, 0])

        self.X = torch.FloatTensor(np.array(self.X))
        self.y = torch.FloatTensor(np.array(self.y))

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


# PyTorch LSTM模型（增强架构）
class EnergyLSTM(nn.Module):
    def __init__(self, input_size, hidden_size=128, num_layers=2, output_size=96):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=0.2
        )

        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, output_size)
        )

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)

        out, _ = self.lstm(x, (h0, c0))
        out = out[:, -1, :]  # 取最后时间步输出
        return self.fc(out)


# 训练流程
def train_model(model, train_loader, val_loader, num_epochs=1):
    criterion = nn.MSELoss()
    optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=3)

    best_loss = float('inf')
    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss = 0
        for X_batch, y_batch in train_loader:
            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            train_loss += loss.item()

        # 验证阶段
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for X_val, y_val in val_loader:
                outputs = model(X_val)
                val_loss += criterion(outputs, y_val).item()

        # 学习率调度
        avg_val_loss = val_loss / len(val_loader)
        scheduler.step(avg_val_loss)

        # 保存最佳模型
        if avg_val_loss < best_loss:
            best_loss = avg_val_loss
            torch.save(model.state_dict(), 'best_model.pth')

        print(f'Epoch {epoch + 1}: Train Loss {train_loss / len(train_loader):.4f}, Val Loss {avg_val_loss:.4f}')


# 主程序
if __name__ == "__main__":
    # 数据准备
    data = generate_15min_data(days=60)
    processed_data = preprocess_data(data)
    features = ['price', 'demand', 'temperature', 'renewable_ratio', 'is_weekend',
                '15min_block', 'day_of_week', 'month', 'price_lag_96',
                'price_lag_192', 'price_lag_288', 'price_7day_ma']

    dataset = EnergyDataset(processed_data[features].values)
    train_size = int(0.8 * len(dataset))
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, len(dataset) - train_size])

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

    # 模型初始化
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = EnergyLSTM(input_size=len(features)).to(device)

    # 训练
    train_model(model, train_loader, val_loader)

    # 加载最佳模型
    model.load_state_dict(torch.load('best_model.pth'))
    model.eval()

    # 预测
    with torch.no_grad():
        last_sequence = dataset.X[-1].unsqueeze(0).to(device)
        forecast = model(last_sequence).cpu().numpy()

    # 逆标准化
    dummy = np.zeros((len(forecast[0]), len(features)))
    dummy[:, 0] = forecast[0]
    forecast_price = dataset.scaler.inverse_transform(dummy)[:, 0]

    # 可视化（保持与之前相同）
    # 生成预测时间戳
    last_timestamp = processed_data['datetime'].iloc[-1]
    future_dates = pd.date_range(start=last_timestamp + pd.Timedelta(minutes=15),
                                 periods=96,
                                 freq="15T")

    plt.figure(figsize=(16, 6))
    plt.plot(processed_data['datetime'][-672:], processed_data['price'][-672:], label='Historical')
    plt.plot(future_dates, forecast_price, label='Forecast', color='orange')
    plt.title("15-Minute Electricity Price Forecast (Next 24 Hours)")
    plt.xlabel("Datetime")
    plt.ylabel("Price ($/MWh)")
    plt.legend()
    plt.grid(True)
    plt.show()